# 新增依赖库（需安装）
# pip install diffusers transformers accelerate torch scipy

import torch
from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
from PIL import Image
import numpy as np

# 1. Stable Diffusion文生图模块
class TextToImageGenerator:
    def __init__(self, model_name="stabilityai/stable-diffusion-2-1"):
        self.pipe = StableDiffusionPipeline.from_pretrained(
            model_name,
            torch_dtype=torch.float16,
            safety_checker=None  # 禁用安全过滤器（可选）
        )
        self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
        self.pipe = self.pipe.to("cuda" if torch.cuda.is_available() else "cpu")
        self.pipe.enable_attention_slicing()  # 减少显存占用:cite[6]
    
    def generate_image(self, prompt, negative_prompt="", steps=28, guidance_scale=7.5, width=512, height=512):
        """ 根据文本生成图像:cite[6]:cite[8] 后续实现做成接口，前端请求时，能将文和图都返回给前端 """
        image = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            width=width,
            height=height
        ).images[0]
        return image

# 2. 提示词优化器（结合语音识别和分词结果）
def optimize_prompt(transcribed_text):
    """ 使用结巴分词提取关键词并构建艺术化提示词:cite[7]:cite[9] """
    import jieba.analyse
    # 提取TF-IDF权重最高的5个关键词
    keywords = jieba.analyse.extract_tags(transcribed_text, topK=5, withWeight=False)
    
    # 构建专业级提示词模板:cite[6]
    quality_tags = "masterpiece, best quality, 4k, detailed"
    prompt = f"{quality_tags}, {', '.join(keywords)}"
    
    # 添加负面提示词防止常见缺陷:cite[8]
    negative_prompt = "worst quality, low quality, deformed, mutation, blurry"
    return prompt, negative_prompt

# 3. 图像后处理模块（可选）
def apply_face_swap(image, target_face_path):
    """ 使用Roop插件实现一键换脸:cite[10] """
    # 需额外安装roop (https://github.com/s0md3v/roop)
    from roop import core
    return core.process(image, target_face_path)

# 4. 主流程整合
if __name__ == "__main__":
    # 原有语音识别流程
    audio_file = record_audio()  # 录音
    text = transcribe_audio(audio_file)  # Whisper识别
    analyze_text(text)  # 结巴分词分析
    
    # 新增图像生成流程
    generator = TextToImageGenerator()
    
    # 优化提示词
    prompt, negative_prompt = optimize_prompt(text)
    print(f"优化后的提示词: {prompt}")
    print(f"负面提示词: {negative_prompt}")
    
    # 生成图像
    image = generator.generate_image(
        prompt,
        negative_prompt=negative_prompt,
        steps=28,          # 推荐采样步数:cite[6]
        guidance_scale=7.5  # 提示词相关性系数:cite[8]
    )
    
    # 保存结果
    image.save("generated_image.png")
    print("图像已保存为 generated_image.png")
    
    # 可选：换脸功能
    if input("是否需要换脸？(y/n): ").lower() == 'y':
        target_face = "target_face.jpg"  # 目标人脸路径
        swapped_image = apply_face_swap(image, target_face)
        swapped_image.save("swapped_image.png")